package cn.actoncode.boot.module.ai.service.chatmodel.impl;

import cn.actoncode.boot.framework.minimax.config.MiniMaxProperties;
import cn.actoncode.boot.framework.minimax.core.bean.MiniMaxRequest;
import cn.actoncode.boot.framework.minimax.core.bean.MiniMaxResponse;
import cn.actoncode.boot.framework.minimax.core.bean.MiniMaxStreamResponse;
import cn.actoncode.boot.framework.minimax.core.client.MiniMaxChatClient;
import cn.actoncode.boot.framework.minimax.core.enums.MiniMaxObjectEnum;
import cn.actoncode.boot.framework.minimax.core.enums.MiniMaxRoleEnum;
import cn.actoncode.boot.module.ai.service.chatmodel.dto.ChatHistoryDTO;
import cn.actoncode.boot.module.ai.service.chatmodel.dto.ChatRespDTO;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;

import static cn.actoncode.boot.framework.common.exception.enums.GlobalErrorCodeConstants.NLP_ERROR;
import static cn.actoncode.boot.framework.common.exception.util.ServiceExceptionUtil.exception;

@Slf4j
@Component("MiniMaxV2")
public class MiniMaxChatServiceImplImpl extends AbstractChatServiceImpl {

    @Resource
    private MiniMaxChatClient miniMaxChatClient;

    @Resource
    private MiniMaxProperties miniMaxProperties;

    @Override
    protected ChatRespDTO doChat(String prompt, List<ChatHistoryDTO> historyList, String system,
                                 Boolean stream, String userId, Long sk) {
        // 构建请求消息体
        List<MiniMaxRequest.Message> messages = new ArrayList<>();
        if (StrUtil.isNotEmpty(system)) {
            MiniMaxRequest.Message message = new MiniMaxRequest.Message();
            message.setRole(MiniMaxRoleEnum.system.name());
            message.setContent(system);
            messages.add(message);
        }
        if (CollUtil.isNotEmpty(historyList)) {
            historyList.forEach(history -> {
                MiniMaxRequest.Message message = new MiniMaxRequest.Message();
                message.setRole(MiniMaxRoleEnum.user.name());
                message.setContent(history.getQuery());
                MiniMaxRequest.Message message1 = new MiniMaxRequest.Message();
                message1.setRole(MiniMaxRoleEnum.assistant.name());
                message1.setContent(history.getContent());
                messages.add(message);
                messages.add(message1);
            });
        }
        // 构建请求体
        MiniMaxRequest miniMaxRequest = miniMaxChatClient.buildRequest("abab5.5-chat", prompt, stream, messages);
        if (stream) {
            // 获取流
            InputStream inputStream = miniMaxChatClient.streamChat(miniMaxRequest, miniMaxProperties.getChatCompletionV2ApiEndpoint());
            // 发送消息
            MiniMaxStreamResponse miniMaxStreamResponse = sendMessage(sk, inputStream);

            StringBuilder sb = new StringBuilder();
            for (MiniMaxStreamResponse.Choices choices: miniMaxStreamResponse.getChoices()) {
                sb.append(choices.getMessage().getContent());
            }

            return ChatRespDTO.builder().prompt(prompt).text(sb.toString())
                    .requestBody(JSON.toJSONString(miniMaxRequest))
                    .totalTokens(miniMaxStreamResponse.getUsage().getTotalTokens())
                    .isMaxLen(false).build();
        } else {
            // 请求
            MiniMaxResponse miniMaxResponse = miniMaxChatClient.chat(miniMaxRequest, miniMaxProperties.getChatCompletionV2ApiEndpoint());
            // 循环拼接结果文本
            StringBuilder sb = new StringBuilder();
            miniMaxResponse.getChoices().forEach(choice -> {
                String content = choice.getMessage().getContent();
                sb.append(content);
            });

            return ChatRespDTO.builder().prompt(prompt).text(sb.toString())
                    .requestBody(JSON.toJSONString(miniMaxRequest))
                    .totalTokens(miniMaxResponse.getUsage().getTotalTokens())
                    .isMaxLen(false).build();
        }
    }

    @Override
    protected MiniMaxStreamResponse sendMessage(Long sk, InputStream inputStream) {
        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream));
        String line;
        MiniMaxStreamResponse response = new MiniMaxStreamResponse();
        try {
            while ((line = reader.readLine()) != null) {
                if(line.length() < 6){
                    continue;
                }
                System.out.println(line);
                String data;
                if (line.startsWith("data:")) {
                    data = line.substring(6);
                } else {
                    data = line;
                }
                MiniMaxStreamResponse streamResponse = JSONObject.parseObject(data, MiniMaxStreamResponse.class);
                if (streamResponse.getBaseResp() != null && streamResponse.getBaseResp().getCode() != 0) {
                    log.info("MiniMaxChat Stream Error : {}", streamResponse.getBaseResp().getMsg());
                    throw exception(NLP_ERROR, "当前功能使用人数过多，请稍后再试");
                } else {
                    if (streamResponse.getObject().equals(MiniMaxObjectEnum.stream.getValue())) {
                        streamResponse.getChoices().forEach(choices -> {
                            super.sendMessage(sk, choices.getDelta().getContent());
                        });
                    } else {
                        response = streamResponse;
                    }
                }
            }
            super.sendMessage(sk, "WRITE_FINISH");
            reader.close();
            inputStream.close();
        } catch (IOException e) {
            log.info(e.getMessage());
            throw exception(NLP_ERROR, "当前功能使用人数过多，请稍后再试");
        }
        return response;
    }
}
